import os
from os import path
import random
import shutil
import yaml
from extrinsics import build_extrinsics

SRC_DATASET = {
    'calib': '/home/data/1354/calib',
    'denorm': '/home/data/1354/denorm',
    'extrinsics': '/home/data/1354/extrinsics',
    'image': '/home/data/1354/image',
    'label': '/home/data/1354/label',
    'sample': '/home/data/1354/sample.txt',
    'depth': '/home/data/1468/depth',
}

DST_BASEDIR = '/project/dataset'

DST_DATASET = {
    'training': path.join(DST_BASEDIR, 'training'),
    'image': path.join(DST_BASEDIR, 'training', 'image_2'),
    'label': path.join(DST_BASEDIR, 'training', 'label_2'),
    'calib': path.join(DST_BASEDIR, 'training', 'calib'),
    'ImageSets': path.join(DST_BASEDIR, 'ImageSets'),
}

def create_dir(dirpath):
    pl = path.split(dirpath)
    for i in range(1, len(pl) + 1):
        p = path.join(*pl[:i])
        if not (path.exists(p) and path.isdir(p)):
            os.mkdir(p)

TRAIN_VAL_SPLIT = 0.8

def load_yaml(filepath):
    with open(filepath, 'r') as f:
        return yaml.safe_load(f)

def write_dataline(name):
    shutil.copy(path.join(SRC_DATASET['image'], f'{name}.jpg'), DST_DATASET['image'])
    shutil.copy(path.join(SRC_DATASET['label'], f'{name}.txt'), DST_DATASET['label'])

    with open(path.join(SRC_DATASET['calib'], f'{name}.txt'), 'r') as f:
        P2 = f.readline()[4:]

    with open(path.join(SRC_DATASET['extrinsics'], f'{name}.yaml'), 'r') as f:
        extrinsics_data = yaml.safe_load(f)
        r, t = extrinsics_data['transform']['rotation'], extrinsics_data['transform']['translation']
        e = build_extrinsics([r['w'], r['x'], r['y'], r['z']], [t['x'], t['y'], t['z']])
        R0_rect = ' '.join([f'{v:.12e}' for l in e[:-1] for v in l])

    origin_3x4 = [(1 if i == j else 0) for i in range(3) for j in range(4)]
    origin_str = ' '.join([f'{v:.12e}' for v in origin_3x4])

    with open(path.join(DST_DATASET['calib'], f'{name}.txt'), 'w') as f:
        print(f'P0: {P2}', file=f)
        print(f'P1: {P2}', file=f)
        print(f'P2: {P2}', file=f)
        print(f'P3: {P2}', file=f)
        print(f'R0_rect: {R0_rect}', file=f)
        print(f'Tr_velo_to_cam: {origin_str}', file=f)
        print(f'Tr_imu_to_velo: {origin_str}', file=f)

if __name__ == '__main__':
    for k, v in DST_DATASET.items():
        create_dir(v)

    with open(SRC_DATASET['sample'], 'r') as f:
        datalines = [s.strip() for s in f.readlines()]
        datalines.sort()

    random.shuffle(datalines)
    # print(datalines[:5])

    split_point = int(TRAIN_VAL_SPLIT * len(datalines))
    train_list = datalines[:split_point]
    val_list = datalines[split_point:]

    # write_dataline(datalines[0])
    for l in datalines:
        write_dataline(l)

    with open(path.join(DST_DATASET['ImageSets'], 'train.txt'), 'w') as f:
        [print(l, file=f) for l in train_list]


    with open(path.join(DST_DATASET['ImageSets'], 'val.txt'), 'w') as f:
        [print(l, file=f) for l in val_list]